import os
import random
import sys
import threading
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed

import pngquant
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

if not os.path.exists("pngquant.exe"):
    print("无法找到pngquant")
    input("按任意键结束")
    sys.exit()
PngquantExe = os.path.join(os.getcwd(), "pngquant.exe")
# 设置 pngquant.exe 文件所在的完整文件名，压缩比例
print("已找到pngquant")
pngquant.config(PngquantExe, min_quality=50, max_quality=60)

if not os.path.exists("jpegtran.exe"):
    print("未找到jpegtran")
    input("按任意键结束")
    sys.exit()
else:
    print("已找到jpegtran")

if not os.path.exists("mozcjpeg.exe"):
    print("未找到mozcjpeg")
    input("按任意键结束")
    sys.exit()
else:
    print("已找到mozcjpeg")


def pd_img(infile):  # 判断图片格式，imghdr不准确
    try:
        im = Image.open(infile, "r")
    except BaseException:
        return None
    return im.format


ysquality1 = 65  # jpegtran质量，实际上是PIL转换质量，jpegtran是无损
ysquality2 = 70  # mozjpeg质量
smooth = 25  # 平滑度

min_ys = 200  # 小于该值图片不压缩，单位kb
maxjd_ys = 1000  # 大于该值图片压缩狠一点，单位kb
maxjd_ys2 = 2000  # 大于该值图片压缩再狠一点，单位kb
after_ys = "ysy"

max_Resolution = 2560  # 分辨率太大的话把图片缩小
zoom_ratio = 0.7  # 缩放比例

tpgs = 'png,jpg,jpeg,JPEG,JPG,PNG,bmp,BMP'

mutex = threading.Lock()  # 文件互斥锁


def pqys(infile, code=0, sccode=1):  # 使用pngquant压缩
    pd1 = 0
    newfile_name = ''
    if not os.path.exists(infile):  # 判断图片还在不在
        return 1
    if os.path.getsize(infile) / 1024 < min_ys:
        return 0
    if pd_img(infile) not in 'png,PNG' and pd_img(infile) is not None and pd_img(infile) != '':  # 不是png就调用jpeg压缩
        print("不是png文件")
        pd1 = pilys(infile=infile, code=code, sccode=sccode)
        if pd1 == 1:
            return 1
        else:
            return 0
    if code == 1:
        newfile_name = os.path.splitext(infile)[0] + ".png"
    elif code == 0 or code == 2:
        newfile_name = os.path.splitext(infile)[0] + "_p" + ".png"
    pngquant.quant_image(infile, dst=newfile_name)
    if not os.path.exists(newfile_name):
        print("图片未压缩成功")
        write_erro(infile)
        return 1
    pqdx = format(os.path.getsize(newfile_name) / 1024, '.1f')
    if code == 2:
        if os.path.exists(newfile_name):  # 如果文件存在
            os.remove(newfile_name)
    if sccode == 0:
        return float(pqdx)
    elif sccode == 1:
        write_scinform(infile, after_ys)
        return 0

# code 1正常使用，会覆盖原始文件
#      2测试使用，会删除测试文件
#      0测试使用，不会删除测试文件


def webpzh(infile, code=0, sccode=1):  # 转换成webp
    if not os.path.exists(infile):  # 判断图片还在不在
        return 1
    if pd_img(infile) is not None and pd_img(infile) in tpgs and pd_img(infile) != '':
        pass
    else:
        return 1  # 不是图片不压缩
    newfile_name = ''
    if code == 1:
        newfile_name = os.path.splitext(infile)[0] + ".webp"
    elif code == 0 or code == 2:
        newfile_name = os.path.splitext(infile)[0] + "_w" + ".webp"
    im = Image.open(infile)
    im.save(newfile_name, "WEBP")
    if not os.path.exists(newfile_name):
        print("图片未压缩成功")
        write_erro(infile)
        return 1
    wedx = format(os.path.getsize(newfile_name) / 1024, '.1f')
    if code == 2:
        if os.path.exists(newfile_name):  # 如果文件存在
            os.remove(newfile_name)
    elif code == 1:
        if os.path.exists(infile):  # 如果文件存在
            os.remove(infile)
    if sccode == 0:
        return float(wedx)
    elif sccode == 1:
        write_scinform(infile, after_ys)
        return 0

# sccode=0 返回文件大小
# sccode=1 返回成功值0


def pilys(infile, ysquality=ysquality1, code=0, sccode=1):  # jpg压缩
    newfile_name1 = ""
    if not os.path.exists(infile):  # 判断图片还在不在
        return 1
    if pd_img(infile) is not None and pd_img(infile) in tpgs and pd_img(infile) != '':
        pass
    else:
        return 1  # 不是图片不压缩
    if os.path.getsize(infile) / 1024 < min_ys:
        return 0
    if code == 1:
        newfile_name1 = os.path.splitext(infile)[0] + ".jpeg"
    else:
        newfile_name1 = os.path.splitext(infile)[0] + "_pl" + ".jpeg"
    im = Image.open(infile)
    im.save(newfile_name1, "JPEG", quality=ysquality)
    if os.path.exists(newfile_name1):
        cmd = "jpegtran -optimize -progressive -outfile " + "\"" + newfile_name1 + "\"" + " " + "\"" + newfile_name1 + "\""
        if os.system(cmd) != 0:
            print("jpegtran压缩失败")
            write_erro(infile)
            return 1
    pldx = format(os.path.getsize(newfile_name1) / 1024, '.1f')
    if code == 2:
        if os.path.exists(newfile_name1):  # 如果文件存在
            os.remove(newfile_name1)
    elif code == 1:
        if os.path.splitext(infile)[1].replace(".", "") != 'jpeg':  # 如果图片不是以jpg结尾
            if os.path.exists(infile):  # 如果文件存在
                os.remove(infile)
    if sccode == 0:
        return float(pldx)
    elif sccode == 1:
        write_scinform(infile, after_ys)
        return 0


def pilys2(infile, code=0, ysquality=ysquality2, sccode=1, ig_code=1):  # jpg模式2压缩
    if not os.path.exists(infile):  # 判断图片还在不在
        print("图片消失了")
        return 0

    if "pl2temp" in os.path.splitext(infile)[0]:  # 如果是临时图片不压缩还要删除
        try:
            os.remove(infile)
        except BaseException:
            print("临时文件无法删除")
            return 1
        return 0

    if os.path.getsize(infile) / 1024 < min_ys:  # 太小了不压缩
        return 0

    newfile_name1 = ""
    newfile_name2 = os.path.splitext(infile)[0] + "_pl2temp" + ".jpeg"  # 临时文件
    newfile_name3 = os.path.splitext(infile)[0] + "_pl2temp2" + ".jpeg"  # 临时文件2
    if code == 1:
        newfile_name1 = os.path.splitext(infile)[0] + ".jpeg"
    else:
        newfile_name1 = os.path.splitext(infile)[0] + "_pl2" + ".jpeg"
    pdimgvalue = pd_img(infile)
    if pdimgvalue is not None and pdimgvalue in tpgs and pdimgvalue != '':
        pass
    else:
        if ig_code == 0:  # 忽视代码，0为忽视，1为不忽视
            return 0
        print(infile + "不是支持的图片")
        return 1  # 不是图片不压缩

    if os.path.exists(newfile_name2):
        try:
            os.remove(newfile_name2)
        except BaseException:
            print("临时文件无法删除")
            return 1
    if os.path.exists(newfile_name3):
        try:
            os.remove(newfile_name3)
        except BaseException:
            print("临时文件无法删除")
            return 1

    im = Image.open(infile)
    if pdimgvalue in 'PNG,png':
        im = im.convert('RGB')
    else:
        im = im

    width, height = im.size  # 图片的长和宽
    if width >= max_Resolution or height >= max_Resolution:  # 图片太大了的话缩小尺寸
        im = im.resize(((int)(width * zoom_ratio), (int)(height * zoom_ratio)), Image.ANTIALIAS)
    im.save(newfile_name3, "JPEG", quality=98)

    if os.path.getsize(newfile_name3) / 1024 < min_ys:  # 太小了不压缩
        if code == 1:
            if os.path.exists(infile):  # 如果文件存在
                try:
                    os.remove(infile)  # 删除原文件
                    if os.path.exists(newfile_name1):
                        os.remove(newfile_name1)
                    if not os.path.exists(infile) and os.path.exists(newfile_name3):
                        os.rename(newfile_name3, newfile_name1)  # 重命名新文件为原文件名
                        write_scinform(newfile_name1, after_ys)
                except BaseException:
                    print("删除原文件错误")
                    return 1
            else:
                if os.path.exists(newfile_name1):
                    os.remove(newfile_name1)
                if not os.path.exists(infile) and os.path.exists(newfile_name3):
                    os.rename(newfile_name3, newfile_name1)  # 重命名新文件为原文件名
                    write_scinform(newfile_name1, after_ys)
        else:
            if os.path.exists(newfile_name3):
                os.rename(newfile_name3, newfile_name1)  # 重命名新文件为原文件名
                write_scinform(newfile_name1, after_ys)
        if os.path.exists(newfile_name2):
            try:
                os.remove(newfile_name2)
            except BaseException:
                print("临时文件无法删除")
        return 0
    if os.path.getsize(newfile_name3) / 1024 > maxjd_ys2:  # 图片太大的话降低质量
        ysquality = int(ysquality / 1.7)
    elif os.path.getsize(newfile_name3) / 1024 > maxjd_ys:
        ysquality = int(ysquality / 1.4)

    cmd = "mozcjpeg -tune-ssim -quality " + str(ysquality) + " -smooth " + str(smooth) + " " + "< \"" + newfile_name3 + "\" >" + " " + "\"" + newfile_name2 + "\""
    if os.system(cmd) != 0:
        print(infile + "mozcjpeg压缩失败")
        if os.path.exists(newfile_name2):
            os.remove(newfile_name2)
        if os.path.exists(newfile_name3):
            os.remove(newfile_name3)
        if ig_code == 0:  # 忽视代码，0为忽视，1为不忽视
            return 0
        write_erro(infile)
        return 1

    if os.path.exists(newfile_name3):
        os.remove(newfile_name3)
    pldx = format(os.path.getsize(newfile_name2) / 1024, '.1f')

    if code == 2 and os.path.exists(newfile_name2):
        os.remove(newfile_name2)
    elif code == 1:
        if os.path.exists(infile):  # 如果文件存在
            try:
                os.remove(infile)  # 删除原文件
                if os.path.exists(newfile_name1) and os.path.exists(newfile_name2):
                    os.remove(newfile_name1)
                if not os.path.exists(infile) and os.path.exists(newfile_name2):
                    os.rename(newfile_name2, newfile_name1)  # 重命名新文件为原文件名
            except BaseException:
                if os.path.exists(newfile_name2):
                    try:
                        os.remove(newfile_name2)
                    except BaseException:
                        print("临时文件无法删除")
                if os.path.exists(newfile_name3):
                    try:
                        os.remove(newfile_name3)
                    except BaseException:
                        print("临时文件无法删除")
                print("删除原文件错误")
                return 1

    if sccode == 0:
        return float(pldx)
    elif sccode == 1:
        write_scinform(newfile_name1, after_ys)
        # print("提示")
        return 0


def pd_size(infile):
    dx = [5000000, 5000000, 5000000, 5000000]
    if pd_img(infile) == 'PNG':
        dx[0] = pqys(infile, code=2, sccode=0)
        dx[1] = webpzh(infile, code=2, sccode=0)
        dx[2] = pilys(infile, code=2, sccode=0)
        dx[3] = pilys2(infile, code=2, sccode=0)
    elif pd_img(infile) == 'JPG' or pd_img(infile) == 'JPEG':
        dx[1] = webpzh(infile, code=2, sccode=0)
        dx[2] = pilys(infile, code=2, sccode=0)
        dx[3] = pilys2(infile, code=2, sccode=0)
    return dx.index(min(dx))


def write_erro(mydir):  # 写入错误文件名
    if not os.path.exists(os.path.abspath(mydir)):
        return 1
    errodir = os.path.split(os.path.abspath(mydir))[0] + "/erro"
    # 若文件不存在，系统自动创建。'a'表示可连续写入到文件，保留原内容
    mutex.acquire()
    f = open(errodir, 'a', encoding='utf-8')
    f.write("\n")  # 换行
    f.write(mydir)
    f.close()
    mutex.release()
    return 0


def write_scinform(infile, flag):
    if not os.path.exists(os.path.abspath(infile)):
        return 1
    mutex.acquire()
    f = open(os.path.split(os.path.abspath(infile))[0] + "/" + flag, 'a', encoding='utf-8')
    f.write("\n")  # 换行
    f.write(infile)
    f.close()
    mutex.release()
    return 0


def write_sc(mydir, flag='123'):  # 写入成功标志
    if not os.path.exists(mydir):
        return 1
    f = open(mydir + "/" + flag, 'a')  # 若文件不存在，系统自动创建。'a'表示可连续写入到文件，保留原内容
    f.close()
    return 0


def find_img(mydir, flag):  # 寻找该文件夹内部图片的详细地址（不含子文件夹）避开已压缩图片
    file_list1 = []
    file_list2 = []
    for folderName, subfolders, filenames in os.walk(mydir):
        for filename in filenames:
            if os.path.splitext(filename)[1].replace(".", "") is not None and os.path.splitext(filename)[1].replace(".", "") in tpgs and os.path.splitext(filename)[1].replace(".", "") != '':
                if os.path.dirname(os.path.join(folderName, filename)) == mydir:
                    file_list1.append(os.path.join(folderName, filename))
    if os.path.exists(mydir + '\\' + flag):  # 避开已压缩图片
        with open(mydir + "/" + flag, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip('\n') != '':
                    file_list2.append(line.strip('\n'))
        return list(set(file_list1) ^ set(file_list2))
    return file_list1


def find_flag(mydir, flag='123'):  # 寻找标志
    for folderName, subfolders, filenames in os.walk(mydir):
        for filename in filenames:
            if filename == flag:
                return 1
    return 0


class MyThread1(threading.Thread):  # 判断那种方式压缩线程
    def __init__(self, img_list1, rd):
        threading.Thread.__init__(self)
        self.img_list1 = img_list1
        self.rd = rd

    def run(self):
        self.result = pd_size(self.img_list1[self.rd])

    def get_result(self):
        try:
            return self.result
        except Exception:
            return None


# print("0 智能判断压缩方式  1 采用png压缩（若图片非png则调用jpg压缩）")
# print("2 转换为WEBP       3 采用jpg模式1压缩（jpegtran）")
# print("4 采用jpg模式2压缩（mozcjpeg）")
def dirys(mydir, model=0, code=0, works=5, flag='123'):  # 压缩一个文件夹的图片
    img_list1 = []
    pd1 = 4
    if model == 1:
        pd1 = 0
    elif model == 2:
        pd1 = 1
    elif model == 3:
        pd1 = 2
    elif model == 4:
        pd1 = 3
    # 判断那种方式压缩合适，0为pngquant压缩，1为转换webq，2为jpg压缩，3为jpg模式2
    data = 0
    if not os.path.isdir(mydir):
        return 1
    if find_flag(mydir, flag) == 1:
        print(mydir + "内部图片已压缩")
        return 1
    img_list1 = find_img(mydir, after_ys)
    if len(img_list1) == 0:
        print(mydir + "没有寻找到需要压缩的图片")
        return 1
    elif len(img_list1) > 6 and model == 0:  # 判断那种压缩方式最合适
        result = []
        threads = []
        ra_list = random.sample(range(0, len(img_list1) - 1), 3)
        for i in range(3):
            t = MyThread1(img_list1, ra_list[i - 1])
            t.start()
            threads.append(t)
        for t in threads:
            t.join()  #  一定执行join,等待子进程执行结束，主进程再往下执行  
        for t in threads:
            result.append(t.get_result())
        pd1 = Counter(result).most_common(1)[0][0]
    elif len(img_list1) < 4 and model == 0:
        pd1 = pd_size(img_list1[random.randint(0, len(img_list1) - 1)])
    if pd1 == 0:
        print("将采用pngquant压缩" + str(mydir) + "下文件")
        with ThreadPoolExecutor(max_workers=works) as t:
            obj_list = []
            for list2 in img_list1:
                print(list2)
                obj = t.submit(pqys, infile=list2, code=code, sccode=1)
                obj_list.append(obj)
            for future in as_completed(obj_list):
                data += future.result()
    elif pd1 == 1:
        print("将采用转换webp方式压缩" + str(mydir) + "下文件")
        with ThreadPoolExecutor(max_workers=works) as t:
            obj_list = []
            for list2 in img_list1:
                obj = t.submit(webpzh, infile=list2, code=code, sccode=1)
                obj_list.append(obj)
            for future in as_completed(obj_list):
                data += future.result()
    elif pd1 == 2:
        print("将采用JPEG（jpegtran模式）方式压缩" + str(mydir) + "下文件")
        with ThreadPoolExecutor(max_workers=works) as t:
            obj_list = []
            for list2 in img_list1:
                obj = t.submit(pilys, infile=list2, code=code, sccode=1)
                obj_list.append(obj)
            for future in as_completed(obj_list):
                data += future.result()
    elif pd1 == 3:
        print("将采用JPEG（mozcjpeg模式）压缩" + str(mydir) + "下文件")
        with ThreadPoolExecutor(max_workers=works) as t:
            obj_list = []
            for list2 in img_list1:
                obj = t.submit(pilys2, infile=list2, code=code, sccode=1)
                obj_list.append(obj)
            for future in as_completed(obj_list):
                data += future.result()
    if data == 0 and flag != 'no':
        write_sc(mydir, flag=flag)
    if data == 0 and os.path.exists(mydir + '\\' + after_ys):  # 删除已处理文件目录
        pass
        os.remove(mydir + '\\' + after_ys)
    if data == 0 and os.path.exists(mydir + '\\' + 'erro'):  # 删除错误文件目录
        os.remove(mydir + '\\' + 'erro')
    return data


if __name__ == '__main__':
    pass
    dir1 = r"cs.jpg"
    # print(pilys(dir1))
    # print(pilys(dir1, sccode=0))
    # dirys(r'Z:\新建文件夹 (2)', model=4, code=1, works=5, flag='123')
    im = Image.open(r"017.png")
    if pd_img(r"017.png") == 'png':
        rim = im.convert('RGB')
    else:
        rim = im
    rim.save(r"0171.jpeg", "JPEG")
